#include "ntl_helpers.cpp"

#include <NTL/LLL.h>
#include <NTL/vector.h>

#include <time.h>
#include <algorithm>
#include <queue>

#include <fstream>
#include <sstream>

#define LIST_THRESHOLD basis.NumCols()
#define RANDOM_THRESHOLD basis.NumCols()
#define GAUSS_THRESHOLD basis.NumCols() 

#define RANDOM_SIEVE
//#define GAUSS_SIEVE
//#define LIST_SIEVE

NTL_CLIENT


void random_reduce(vec *v, vector<vec> &list);
vec random_sieve(mat_ZZ basis);

void gauss_reduce(vec *t, vector<vec> &list, queue<vec> &S);
vec gauss_sieve(mat_ZZ basis);
void sorted_add(vec *t, vector<vec> &list);

vec list_sieve(mat_ZZ basis);
void list_reduce(vec *t, vector<vec> &list);


unsigned long random_sieve_time;

unsigned long random_iterations;

unsigned long gauss_sieve_time;
unsigned long gauss_iterations;

unsigned long list_sieve_time;
unsigned long list_iterations;

double pi = atan(1) * 4;

int DIM;

int main(int argc, char** argv)
{
  srand(time(0));

  for(int i = atoi(argv[1]); i <= atoi(argv[2]); i++) 
  {
    for(int j = 0; j < 10; j++)
    {
      DIM = i;
      mat_ZZ basis = randomBasis(i);
      mat_ZZ LLL_basis = basis;
      G_BKZ_FP(basis, 0.99, 20);

      //some useful stats to keep track of
      ZZ shortest_mag_sq; 

      //our streams for output
      ofstream f_random_time;
      ofstream f_random_iterations;
      ofstream f_random_correctness;
      ofstream f_gauss_time;
      ofstream f_gauss_iterations;
      ofstream f_gauss_correctness;
      ofstream f_list_time;
      ofstream f_list_iterations;
      ofstream f_list_correctness;
      string name;

      ZZ det;
      long i = LLL(det, LLL_basis);
      shortest_mag_sq =  mag_sq(LLL_basis[LLL_basis.NumRows() - i]);

#ifdef RANDOM_SIEVE
      vec r = random_sieve(basis);

      //record random_times
      name = "random_time_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_random_time.open(name.c_str(), ios::app);
      f_random_time << random_sieve_time << endl;
      random_sieve_time = 0;

      //record random_iterations
      name = "random_iterations_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_random_iterations.open(name.c_str(), ios::app);
      f_random_iterations << random_iterations << endl;
      random_iterations = 0;

      //record random_correctness
      name = "random_correctness_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_random_correctness.open(name.c_str(), ios::app);
      f_random_correctness << (r.mag_sq <= shortest_mag_sq) << endl;

#endif

#ifdef GAUSS_SIEVE
      vec g = gauss_sieve(basis);

      //record gauss_times
      name = "gauss_time_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_gauss_time.open(name.c_str(), ios::app);
      f_gauss_time << gauss_sieve_time << endl;
      gauss_sieve_time = 0;

      //record gauss_iterations
      name = "gauss_iterations_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_gauss_iterations.open(name.c_str(), ios::app);
      f_gauss_iterations << gauss_iterations << endl;
      gauss_iterations = 0;

      //record gauss_correctness
      name = "gauss_correctness_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_gauss_correctness.open(name.c_str(), ios::app);
      f_gauss_correctness << (g.mag_sq <= shortest_mag_sq) << endl;
#endif

#ifdef LIST_SIEVE
      vec l = list_sieve(basis);

      //record list sieve time
      name = "list_time_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_list_time.open(name.c_str(), ios::app);
      f_list_time << list_sieve_time << endl;
      list_sieve_time = 0;

      //record list sieve iterations
      name = "list_iterations_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_list_iterations.open(name.c_str(), ios::app);
      f_list_iterations << list_iterations << endl;
      list_iterations = 0;

      //record list correctness
      name = "list_correctness_";
      {
        stringstream convert;
        convert << i;
        name.append(convert.str());
      }
      f_list_correctness.open(name.c_str(), ios::app);
      f_list_correctness << (l.mag_sq <= shortest_mag_sq) << endl;
#endif

    }
  }

  return 0;
}


/*
 * Finds the shortest vector in a lattice by taking a random vector from the
 * lattice, using only 0 1's, and then reducing it using random vectors from
 * our list, which we made from previous vectors we reduced. 
 * We keep doing this until the list does not grow a number of times, and 
 * finally we return the shortest vector we found
 *
 * This uses RANDOM_THRESHOLD to set the threshold
 */
vec random_sieve(mat_ZZ basis)
{
  clock_t begin = clock();
  vector<vec> list;
  vec t;
  vec_vec_ZZ listend;
  mat_ZZ newBasis;
  ZZ basis_det; //determinant of basis
  int collisions = 0;
  int dimension = basis.NumCols();
  int max_size;

  determinant(basis_det, basis, 1);
  while(collisions < RANDOM_THRESHOLD)
  {
    random_iterations++;
    int contains = 0;
    //get random lattice vector
    t = randomVec(basis);
    //reduce v
    random_reduce(&t, list);

    //if it ends up being nonzero we want to add it

    if(t.mag_sq > 0)
    {
      //check if it's already in list
      for(int i = 0; i < list.size(); i++)
      {
        if(list[i].v == t.v)
        {
          contains = 1;
          break;
        }
      }
      if(!contains)
      {
        sorted_add(&t, list); //put it in list
        collisions = 0; //we added one so we restart our collision counter

        if(list.size() > max_size)
          max_size = list.size();
      }
      else //if it's nonzero but we already have it..
        collisions++;

    }
    //otherwise if it's zero, it's a collision
    else
      collisions++;
    //or, if we already have it, collisions++
    if(contains)
      collisions++;
  }


  //now we just find the shortest one
  if(list.size() > 0)
  {
    t = list[0];
    for(int i = 1; i < list.size(); i++)
    {
      if(list[i].mag_sq < t.mag_sq)
        t = list[i];
    }
  }
  else
  {
    setVec(&t, basis[0]);
    for(int i = 1; i < basis.NumRows(); i++)
      if(mag_sq(basis[i]) < t.mag_sq)
        setVec(&t, basis[i]);
  }

  random_sieve_time += (clock() - begin);

  return t;
}

/* 
 * reduces t using a randomly chosen vector in list. We repeat until we have
 * tried all vectors in the list shorter than t. We assume that list is sorted
 */
void random_reduce(vec *t, vector<vec> &list)
{
  clock_t begin = clock();

  vec red;
  int i; //random pick from choices

  RR delta_sq;
  delta_sq = 1 - 1/(double)t->v.length();
  delta_sq = sqr(delta_sq);

  int long_ind; //this is first index that is too long in list
  for(long_ind = 0; (long_ind < list.size()) && 
      (list[long_ind].mag_sq < t->mag_sq); long_ind++) 
    ;

  //keeps track of what we can still pick
  vector<int> indices;
  for(int i = 0; i < long_ind; i++)
    indices.push_back(i);

  while(indices.size() > 0) //while there are still things to try
  {
    i = rand()%indices.size(); //random choice from indices
    red = reduce(t, &list[indices[i]]); //reduce by the randomly chosen vec
    //remove the one we picked from indices so we don't repeat
    indices.erase(indices.begin() + i);

    //if it gets reduced
    if(MakeRR(red.mag_sq, 0) < delta_sq * MakeRR(t->mag_sq, 0))
    {
      *t = red; //we should reduce v
      //reset picked_ind and indices
      indices.clear();
      for(long_ind = 0; (long_ind < list.size()) && 
          (list[long_ind].mag_sq < t->mag_sq); long_ind++)
        ;
      for(int i = 0; i < long_ind; i++)
        indices.push_back(i);
    }
  }
}

/*
 * Gauss Sieve as outlined in paper! Shortens a random vector as much as
 * possible, then tries reducing the other vectors and if the others one can
 * be reduced, we do the same for those.
 */
vec gauss_sieve(mat_ZZ basis)
{
  clock_t begin = clock();

  vector<vec> list;
  queue<vec> S;
  int collisions = 0;
  vec t;
  while(collisions < GAUSS_THRESHOLD)
  {
    gauss_iterations++;
    if(!S.empty() )
    {
      t = S.front();
      S.pop();
    }
    else
      t = randomVec(basis);
    gauss_reduce(&t, list, S);
    if(t.mag_sq == 0)
      collisions++;
    else
      sorted_add(&t, list);
  }

  gauss_sieve_time += clock() - begin;

  return list[0];
}


/*
 * gauss_sieve that only grows longer and doesn't try reducing vectors already
 * in the list
 */
vec list_sieve(mat_ZZ basis)
{
  clock_t begin = clock();

  vector<vec> list;
  int collisions = 0;
  vec t;
  while(collisions < LIST_THRESHOLD)
  {
    list_iterations++;
    t = randomVec(basis);
    list_reduce(&t, list);
    if(t.mag_sq == 0)
      collisions++;
    else
      sorted_add(&t, list);
  }

  list_sieve_time += clock() - begin;
  return list[0];
}

/* 
 * reduces t using vectors in the list, which we assume to be sorted
 */
void list_reduce(vec *t, vector<vec> &list)
{
  for(int i = 0; i < list.size(); i++)
  {
    //if the list vector is longer we are done!
    if(list[i].mag_sq > t->mag_sq)
      break;
    vec red = reduce(t, &list[i]);
    if(red.mag_sq < t->mag_sq)
    {
      *t = red;
      i = 0;
    }
  }
}


/*
 * reduces t using vectors in the list as much as possible, and then tries 
 * reducing the vectors already in the list by t, and if they get reduced
 * we add them to S. 
 * For optimization reasons we stop when we reach a vector longer than t
 */
void gauss_reduce(vec *t, vector<vec> &list, queue<vec> &S)
{
  clock_t begin = clock();

  for(int i = 0; i < list.size(); i++)
  {
    //if it's longer we are done
    if(list[i].mag_sq > t->mag_sq)
      break;

    //otherwise, we try reducing it and see if it gets smaller
    vec red = reduce(t, &list[i]);
    if(red.mag_sq < t->mag_sq)
    {
      *t = red;
      i = 0;
      //if we reduce it, start over!
    }
  }

  //here, for all the vectors which get smaller by using t, we add to queue S
  for(int i = 0; i < list.size(); i++)
  {
    if(list[i].mag_sq > t->mag_sq)
      break;
    vec red = reduce(t, &list[i]);
    if(red.mag_sq < list[i].mag_sq)
    {
      S.push(list[i]);
      list.erase(list.begin() + i);
    }
  }

}


/* 
 * This function needs to add the vec to the list in a way such that the 
 * list remains sorted, and we can assume it is already sorted
 */
void sorted_add(vec *t, vector<vec> &list)
{
  //we start from the top since it's more likely that we will add a short one
  //just increments i till we are at a point that t is shorter than list[i]
  int i = 0;
  while(i < list.size() && list[i].mag_sq < t->mag_sq)
    i++;

  list.insert(list.begin() + i, *t);;
}
